import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge
import cv2
import numpy as np
from ultralytics import YOLO
from std_msgs.msg import String
import json
import message_filters


class Detect3DNode(Node):
    def __init__(self):
        super().__init__('detect_3d')
        self.model = YOLO('yolov8n.pt')
        self.bridge = CvBridge()

        self.color_camera_info = None
        self.depth_camera_info = None

        self.info_color_sub = self.create_subscription(
            CameraInfo,
            '/camera/camera/color/camera_info',
            self.color_info_callback,
            10
        )
        self.info_depth_sub = self.create_subscription(
            CameraInfo,
            '/camera/camera/depth/camera_info',
            self.depth_info_callback,
            10
        )

        self.color_sub = message_filters.Subscriber(self, Image, '/camera/camera/color/image_raw')
        self.depth_sub = message_filters.Subscriber(self, Image, '/camera/camera/aligned_depth_to_color/image_raw')

        self.ts = message_filters.ApproximateTimeSynchronizer(
            [self.color_sub, self.depth_sub],
            queue_size=10,
            slop=0.1
        )
        self.ts.registerCallback(self.sync_callback)

        self.result_pub = self.create_publisher(String, '/detected_objects_3d', 10)

        self.get_logger().info("YOLOv8 + D435i 3D Detection Node Started (Sync Mode).")

    def color_info_callback(self, msg):
        if self.color_camera_info is None:
            self.color_camera_info = msg
            K = msg.k
            self.get_logger().info(f"Color Intrinsics (K): fx={K[0]:.2f}, fy={K[4]:.2f}, cx={K[2]:.2f}, cy={K[5]:.2f}")
            # 收到信息后取消订阅，避免不必要的资源占用
            self.destroy_subscription(self.info_color_sub)

    def depth_info_callback(self, msg):
        if self.depth_camera_info is None:
            self.depth_camera_info = msg
            K = msg.k
            self.get_logger().info(f"Depth Intrinsics (K): fx={K[0]:.2f}, fy={K[4]:.2f}, cx={K[2]:.2f}, cy={K[5]:.2f}")
            # 收到信息后取消订阅
            self.destroy_subscription(self.info_depth_sub)

    def sync_callback(self, color_msg, depth_msg):
        if self.color_camera_info is None:
            self.get_logger().warn("Waiting for color camera info...", throttle_duration_sec=5)
            return

        try:
            color_image = self.bridge.imgmsg_to_cv2(color_msg, "bgr8")

            # 转换深度图并处理单位
            depth_image = self.bridge.imgmsg_to_cv2(depth_msg, "32FC1")
            if depth_msg.encoding == '16UC1':
                depth_image = depth_image / 1000.0  # 毫米 -> 米

        except Exception as e:
            self.get_logger().error(f"Image conversion failed: {e}")
            return

        try:
            results = self.model(color_image, verbose=False)
        except Exception as e:
            self.get_logger().error(f"YOLOv8 inference failed: {e}")
            return

        if not results:
            return

        result = results[0]
        im_array = result.plot()
        boxes = result.boxes

        for box in boxes:
            xywh = box.xywh[0].cpu().numpy()
            cx, cy = int(xywh[0]), int(xywh[1])
            conf = float(box.conf[0])
            cls_id = int(box.cls[0])
            label = result.names[cls_id]

            if not (0 <= cy < depth_image.shape[0] and 0 <= cx < depth_image.shape[1]):
                continue

            # --- 稳健地获取深度值 ---
            patch_size = 5
            half_patch = patch_size // 2

            v1 = max(0, cy - half_patch)
            v2 = min(depth_image.shape[0] - 1, cy + half_patch)
            u1 = max(0, cx - half_patch)
            u2 = min(depth_image.shape[1] - 1, cx + half_patch)

            depth_patch = depth_image[v1:v2 + 1, u1:u2 + 1]

            # 过滤无效值 (nan, 0, 或太近的值)
            valid_depths = depth_patch[np.isfinite(depth_patch) & (depth_patch > 0.1)]

            if len(valid_depths) == 0:
                continue

            z = np.median(valid_depths)
            # --- 深度获取结束 ---

            if z > 3.0:  # 限制最大距离
                continue

            K = self.color_camera_info.k
            fx = K[0]
            fy = K[4]
            cx0 = K[2]
            cy0 = K[5]

            x = (cx - cx0) * z / fx
            y = (cy - cy0) * z / fy
            xyz_mm = [round(v * 1000, 2) for v in [x, y, z]]

            msg_data = {
                "class": label,
                "confidence": round(conf, 3),
                "pixel_center": [cx, cy],
                "camera_xyz_mm": xyz_mm
            }
            self.result_pub.publish(String(data=json.dumps(msg_data)))

            text = f"{label} {conf:.2f} | XYZ({xyz_mm[0]:.0f},{xyz_mm[1]:.0f},{xyz_mm[2]:.0f})mm"
            font_scale = 0.5
            thickness = 1
            text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)[0]

            text_x = max(0, min(cx - text_size[0] // 2, color_image.shape[1] - text_size[0]))
            text_y = max(0, min(cy + 20, color_image.shape[0]))

            bg_color = (0, 0, 0)
            alpha = 0.7
            overlay = im_array.copy()
            cv2.rectangle(
                overlay,
                (text_x, text_y - text_size[1] - 5),
                (text_x + text_size[0] + 5, text_y + 5),
                bg_color,
                -1
            )
            cv2.addWeighted(overlay, alpha, im_array, 1 - alpha, 0, im_array)

            cv2.putText(
                im_array,
                text,
                (text_x, text_y),
                cv2.FONT_HERSHEY_SIMPLEX,
                font_scale,
                (255, 255, 255),
                thickness,
                cv2.LINE_AA
            )
            cv2.circle(im_array, (cx, cy), 4, (0, 255, 0), -1)

        cv2.imshow("ROS2 + YOLOv8 + 3D", im_array)
        cv2.waitKey(1)


def main(args=None):
    rclpy.init(args=args)
    node = Detect3DNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()
        cv2.destroyAllWindows()


if __name__ == '__main__':
    main()